package com.tuniu.agents.common.message;

import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.util.I18nUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.context.i18n.LocaleContextHolder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;
@Slf4j
public class FluxSessionManager<T> {
    private final Map<String, Sinks.Many<T>> responseSinks = new ConcurrentHashMap<>();
    private final Map<String, Object> responseLocks = new ConcurrentHashMap<>();
    private final Map<String, Queue<Flux<T>>> responseQueues = new ConcurrentHashMap<>();
    private final Map<String, AtomicBoolean> processingFlags = new ConcurrentHashMap<>();
    private final Set<String> activeResponseIds = ConcurrentHashMap.newKeySet();

    // Main response ID, used to merge all other responses
    private String mainResponseId;
    // Flag to indicate whether the main response ID has been set
    private final AtomicBoolean mainResponseIdSet = new AtomicBoolean(false);
    // Flag to indicate whether the main response has been completed
    private final AtomicBoolean mainResponseCompleted = new AtomicBoolean(false);

    /**
     * Initialize and start a session, returning a Flux that can receive data.
     * Sets the main response ID, and all other responses will be merged into this response.
     */
    public Flux<T> initializeAndStartSession(Supplier<Flux<T>> fluxSupplier, String responseId) {
        // Call the overloaded method with initialMessage as null, indicating no initial message
        return initializeAndStartSession(fluxSupplier, responseId, null);
    }

    /**
     * Initialize and start a session, returning a Flux that can receive data.
     * Sets the main response ID, and all other responses will be merged into this response.
     * This overloaded method allows sending an initial message immediately upon initialization, which is ideal for sending prompts like "Processing...".
     */
    public Flux<T> initializeAndStartSession(Supplier<Flux<T>> fluxSupplier, String responseId, T initialMessage) {
        // Set the main response ID
        setMainResponseId(responseId);

        synchronized (getResponseLock(responseId)) {
            Sinks.Many<T> sink = responseSinks.computeIfAbsent(responseId,
                    id -> Sinks.many().unicast().onBackpressureBuffer());

            responseQueues.computeIfAbsent(responseId, id -> new ConcurrentLinkedQueue<>());
            processingFlags.computeIfAbsent(responseId, id -> new AtomicBoolean(false));
            activeResponseIds.add(responseId);

            // Immediately send the initial message
            if (initialMessage != null) {
                Sinks.EmitResult result = sink.tryEmitNext(initialMessage);
                if (result.isFailure()) {
                    handleEmitFailure(result, responseId, "Initial message");
                }
            }

            // First, create and return sink.asFlux, then asynchronously obtain and process the result of fluxSupplier
            Flux<T> resultFlux = sink.asFlux()
                    .doFinally(signalType -> {
                        // If the main response is completed, mark the main response as completed
                        if (responseId.equals(mainResponseId)) {
                            mainResponseCompleted.set(true);
                            // Clean up all resources
                            cleanupAllResources();
                        } else {
                            // Non-main response, only clean up its own resources
                            cleanupResources(responseId);
                        }
                    });
            Locale locale = LocaleContextHolder.getLocale();
            // Asynchronously obtain and process the result of fluxSupplier
            // Use a separate thread to perform time-consuming operations to avoid blocking the current thread
            reactor.core.scheduler.Schedulers.boundedElastic().schedule(() -> {
                try {
                    LocaleContextHolder.setLocale(locale);
                    // In a separate thread, obtain and start the service with flux
                    Flux<T> flux = fluxSupplier.get();
                    startServiceInBackground(flux, responseId);
                } catch (Exception e) {
                    // Handle any exceptions that may occur
                    log.error("Error occurred while obtaining or processing flux: ", e);
                    emitErrorSafely(responseId, e);
                }
            });

            return resultFlux;
        }
    }

    /**
     * Set the main response ID
     */
    private void setMainResponseId(String responseId) {
        if (mainResponseIdSet.compareAndSet(false, true)) {
            this.mainResponseId = responseId;
        }
    }

    /**
     * Output a single message to the specified response stream
     */
    public void outputFlux(String responseId, T message) {
        // If the main response is completed, do not output any content
        if (mainResponseCompleted.get()) {
            return;
        }

        synchronized (getResponseLock(responseId)) {
            Sinks.Many<T> sink = responseSinks.get(responseId);
            if (sink != null) {
                Sinks.EmitResult result = sink.tryEmitNext(message);
                if (result.isFailure()) {
                    handleEmitFailure(result, responseId, "Single message");
                }
            }
        }
    }

    /**
     * Output a Flux stream to the specified response stream
     * Same responseId streams are processed in order, different responseId streams are processed in parallel
     * If responseId does not exist, it will be automatically created and registered to the main response
     */
    public void outputFlux(String responseId, Flux<T> flux) {
        // If the main response is completed, do not output any content
        if (mainResponseCompleted.get()) {
            return;
        }

        // Ensure the responseId is initialized
        ensureResponseIdInitialized(responseId);

        // Add flux to the queue
        Queue<Flux<T>> queue = responseQueues.get(responseId);
        if (queue != null) {
            queue.offer(flux);
            // Try to process the queue
            processNextInQueue(responseId);
        }
    }

    /**
     * Ensure the responseId is initialized and automatically registered to the main response
     */
    private void ensureResponseIdInitialized(String responseId) {
        if (!activeResponseIds.contains(responseId)) {
            synchronized (getResponseLock(responseId)) {
                if (!activeResponseIds.contains(responseId)) {
                    // Create a new sink
                    Sinks.Many<T> sink = Sinks.many().unicast().onBackpressureBuffer();
                    responseSinks.put(responseId, sink);
                    responseQueues.put(responseId, new ConcurrentLinkedQueue<>());
                    processingFlags.put(responseId, new AtomicBoolean(false));
                    activeResponseIds.add(responseId);

                    // If the main response ID is set, automatically register this new response to the main response
                    if (mainResponseIdSet.get() && !responseId.equals(mainResponseId)) {
                        // Forward the output of the new response to the main response
                        sink.asFlux().subscribe(
                                item -> {
                                    // If the main response is not completed, forward
                                    if (!mainResponseCompleted.get()) {
                                        outputFlux(mainResponseId, item);
                                    }
                                },
                                error -> {},
                                () -> {}
                        );
                    }
                }
            }
        }
    }

    /**
     * Process the next Flux in the queue
     */
    private void processNextInQueue(String responseId) {
        // If the main response is completed, do not process any queues
        if (mainResponseCompleted.get()) {
            return;
        }

        AtomicBoolean isProcessing = processingFlags.get(responseId);
        if (isProcessing == null) return;

        // If already processing, return directly
        if (isProcessing.getAndSet(true)) {
            return;
        }

        Queue<Flux<T>> queue = responseQueues.get(responseId);
        if (queue == null || queue.isEmpty()) {
            isProcessing.set(false);
            return;
        }

        Flux<T> nextFlux = queue.poll();
        if (nextFlux == null) {
            isProcessing.set(false);
            return;
        }

        // Process the current Flux, then process the next in the queue
        nextFlux.subscribe(
                item -> {
                    // If the main response is not completed, output
                    if (!mainResponseCompleted.get()) {
                        emitNextSafely(responseId, item);
                    }
                },
                error -> {
                    // If the main response is not completed, handle error
                    if (!mainResponseCompleted.get()) {
                        emitErrorSafely(responseId, error);
                    }
                    isProcessing.set(false);
                    // If the main response is not completed, handle the next
                    if (!mainResponseCompleted.get()) {
                        processNextInQueue(responseId);
                    }
                },
                () -> {
                    isProcessing.set(false);
                    // If the main response is not completed, handle the next
                    if (!mainResponseCompleted.get()) {
                        processNextInQueue(responseId);
                    }
                }
        );
    }

    /**
     * Safely emit the next element
     */
    private void emitNextSafely(String responseId, T item) {
        // If the main response is completed, do not output any content
        if (mainResponseCompleted.get()) {
            return;
        }

        synchronized (getResponseLock(responseId)) {
            Sinks.Many<T> sink = responseSinks.get(responseId);
            if (sink != null) {
                Sinks.EmitResult result = sink.tryEmitNext(item);
                if (result.isFailure()) {
                    handleEmitFailure(result, responseId, "Stream element");
                }
            }
        }
    }

    /**
     * Safely emit error signal
     */
    private void emitErrorSafely(String responseId, Throwable error) {
        // If the main response is completed, do not output any content
        if (mainResponseCompleted.get()) {
            return;
        }

        synchronized (getResponseLock(responseId)) {
            Sinks.Many<T> sink = responseSinks.get(responseId);
            if (sink != null) {
                Sinks.EmitResult result = sink.tryEmitError(error);
                if (result.isFailure()) {
                    handleEmitFailure(result, responseId, "Error signal");
                }
            }
        }
    }

    /**
     * Handle the failure of emitting
     */
    private void handleEmitFailure(Sinks.EmitResult result, String responseId, String operationType) {
        log.error("Emitting to response ID " + responseId + " " + operationType + " failed: " + result);
    }

    /**
     * Start the service in the background
     */
    private void startServiceInBackground(Flux<T> flux, String responseId) {
        // 60 seconds timeout
        flux.timeout(Duration.ofSeconds(60))
                .onErrorResume(error -> {
                    // Add error handling, convert errors to default messages
                    log.error("FluxSessionManager.startServiceInBackground error: " + error.getMessage());
                    try {
                        // Create a new stream containing error information
                        @SuppressWarnings("unchecked")
                        T errorMessage = (T) BaseAgentNewResp.sendAnswerMsg("", I18nUtil.getMessage(I18nConstants.COMMON_GLOBAL_ERROR_MESSAGE_KEY) );
                        return Flux.just(errorMessage);
                    } catch (ClassCastException e) {
                        log.error("Cannot convert error message to T type: " + e.getMessage());
                        return Flux.empty();
                    }
                }).subscribe(item -> {
                    // When item is null, directly output a string message
                    if (item == null || (item instanceof String && StringUtils.isBlank((String) item))) {
                        // Here, a separate method is needed to output a string because the type may differ from T
                         outputNullItemMessage(responseId, BaseAgentNewResp.sendAnswerMsg("", I18nUtil.getMessage(I18nConstants.COMMON_GLOBAL_ERROR_MESSAGE_KEY)));
                    } else {
                        outputFlux(responseId, item);
                    }
                },
                error -> {
                    log.error("Error occurred while processing stream: " + error.getMessage());
                    emitErrorSafely(responseId, error);

                    // If the main response encounters an error, mark the main response as completed
                    if (responseId.equals(mainResponseId)) {
                        mainResponseCompleted.set(true);
                        // Clean up all resources
                        cleanupAllResources();
                    }
                },
                () -> {
                    // If the main response is completed, mark the main response as completed
                    if (responseId.equals(mainResponseId)) {
                        mainResponseCompleted.set(true);
                        // Complete the main response
                        completeSink(responseId);
                        // Clean up all resources
                        cleanupAllResources();
                    }
                }
        );
    }

    /**
     * When receiving a null item, output a string message
     * Note: This method needs specific implementation depending on your actual T type and output mechanism
     */
    private void outputNullItemMessage(String responseId, String message) {
        // Here, the implementation depends on your T type
        // If T is String, you can try:
        try {
            @SuppressWarnings("unchecked")
            T convertedMessage = (T) message;
            outputFlux(responseId, convertedMessage);
        } catch (ClassCastException e) {
            log.error("Cannot convert string message to T type: " + e.getMessage());
            // Consider other ways to notify the user
        }
    }

    /**
     * Complete the specified sink
     */
    private void completeSink(String responseId) {
        synchronized (getResponseLock(responseId)) {
            Sinks.Many<T> sink = responseSinks.get(responseId);
            if (sink != null) {
                sink.tryEmitComplete();
            }
        }
    }

    /**
     * Get the response lock
     */
    private Object getResponseLock(String responseId) {
        return responseLocks.computeIfAbsent(responseId, key -> new Object());
    }

    /**
     * Clean up resources for the specified response ID
     */
    private void cleanupResources(String responseId) {
        synchronized (getResponseLock(responseId)) {
            responseSinks.remove(responseId);
            responseLocks.remove(responseId);
            responseQueues.remove(responseId);
            processingFlags.remove(responseId);
            activeResponseIds.remove(responseId);
        }
    }

    /**
     * Clean up all resources
     */
    private void cleanupAllResources() {
        // Copy a set of activeResponseIds to avoid concurrent modification exceptions
        Set<String> responseIdsCopy = new HashSet<>(activeResponseIds);
        for (String responseId : responseIdsCopy) {
            // Complete all sinks
            completeSink(responseId);
            // Clean up resources
            cleanupResources(responseId);
        }
    }

    /**
     * Get the managed Flux
     */
    public Flux<T> getManagedFlux(String responseId) {
        ensureResponseIdInitialized(responseId);
        synchronized (getResponseLock(responseId)) {
            Sinks.Many<T> sink = responseSinks.get(responseId);
            if (sink != null) {
                return sink.asFlux()
                        .doFinally(signalType -> {
                            // If the main response is completed, mark the main response as completed
                            if (responseId.equals(mainResponseId)) {
                                mainResponseCompleted.set(true);
                                // Clean up all resources
                                cleanupAllResources();
                            } else {
                                // Non-main response, only clean up its own resources
                                cleanupResources(responseId);
                            }
                        });
            }
            return Flux.empty();
        }
    }

    /**
     * Manually complete the main response
     */
    public void completeMainResponse() {
        if (mainResponseIdSet.get()) {
            mainResponseCompleted.set(true);
            completeSink(mainResponseId);
            cleanupAllResources();
        }
    }
}